import numpy as np

from traits.api import SingletonHasTraits, Trait, Instance, Property, String, Range, Float, Int, Bool, Array, Enum, Button, on_trait_change, cached_property, Code, List, NO_COMPARE
from traitsui.api import View, Item, HGroup, VGroup, VSplit, Tabbed, EnumEditor, TextEditor, Group
from enable.api import Component, ComponentEditor
from chaco.api import ArrayPlotData, Plot, Spectral, PlotLabel

from traitsui.file_dialog import save_file
from traitsui.menu import Action, Menu, MenuBar

import time
import threading
import logging

from tools.emod import ManagedJob

import hardware.api as ha

from analysis import fitting

from tools.utility import GetSetItemsHandler, GetSetItemsMixin

class ODMRHandler( GetSetItemsHandler ):

    def saveLinePlot(self, info):
        filename = save_file(title='Save Line Plot')
        if filename is '':
            return
        else:
            if filename.find('.png')==-1:
                filename=filename+'.png'
            info.object.save_line_plot(filename)

    def saveMatrixPlot(self, info):
        filename = save_file(title='Save Matrix Plot')
        if filename is '':
            return
        else:
            if filename.find('.png')==-1:
                filename=filename+'.png'
            info.object.save_matrix_plot(filename)
    
    def saveAll(self, info):
        filename = save_file(title='Save All')
        if filename is '':
            return
        else:
            info.object.save_all(filename)


class ODMR( ManagedJob, GetSetItemsMixin ):
    """Provides ODMR measurements."""

    keep_data = Bool(False) # helper variable to decide whether to keep existing data

    resubmit_button = Button(label='resubmit', desc='Submits the measurement to the job manager. Tries to keep previously acquired data. Behaves like a normal submit if sequence or time bins have changed since previous run.')    

    power = Range(low=-100., high=25., value=-20, desc='Power [dBm]', label='Power [dBm]', mode='text', auto_set=False, enter_set=True)

    frequency_begin = Range(low=1,      high=20e9, value=2.85e9,    desc='Start Frequency [Hz]',    label='Begin [Hz]')
    frequency_end   = Range(low=1,      high=20e9, value=2.88e9,    desc='Stop Frequency [Hz]',     label='End [Hz]')
    frequency_delta = Range(low=1e-3,   high=20e9, value=1e6,       desc='frequency step [Hz]',     label='Delta [Hz]')

    t_pi   = Range(low=1., high=100000., value=1000., desc='length of pi pulse [ns]', label='pi [ns]', mode='text', auto_set=False, enter_set=True)
    laser = Range(low=1., high=10000., value=300., desc='laser [ns]', label='laser [ns]', mode='text', auto_set=False, enter_set=True)
    wait = Range(low=1., high=10000., value=1000., desc='wait [ns]', label='wait [ns]', mode='text', auto_set=False, enter_set=True)

    pulsed = Bool(False, label='pulsed', enabled_when='state == "idle"')
    sequence = Property(trait=List, depends_on='laser,wait,t_pi')
    
    n_lines = Range (low=1, high=10000, value=50, desc='Number of lines in Matrix', label='Matrix lines', mode='text', auto_set=False, enter_set=True)

    seconds_per_point = Range(low=20e-3, high=1, value=20e-3, desc='Seconds per point', label='Seconds per point', mode='text', auto_set=False, enter_set=True)
    run_time = Range(low=1, high=1000000, value=120, desc='Run Time [s]', label='Run Time [s]', auto_set=False, enter_set=True)
    elapsed_time = Int(value=0, desc='Elapsed Time [s]', label='Elapsed Time [s]', mode='text')
    progress = Int(value=0, desc='Progress [%]', label='Progress [%]', mode='text')

    perform_fit = Bool(False, label='perform fit')

    number_of_dips = Trait( 'auto', String('auto', auto_set=False, enter_set=True), Int(10000., desc='Number of Lorentzians used in fit', label='N', auto_set=False, enter_set=True))

    #fit_parameters = Property( trait=Array(value=(np.NaN, np.NaN, np.NaN, np.NaN)), depends_on='counts,number_of_dips,perform_fit' )
    fit_parameters = Property( trait=Array(value=(0., 0., 0., 0.)), depends_on='counts,number_of_dips,perform_fit' )

    #odmr_frequencies = Property( trait=Array(value=(np.NaN,)), depends_on='fit_parameters', label='frequencies [Hz]' ) 
    odmr_frequencies = Property( trait=Array(value=(0.,)), depends_on='fit_parameters', label='frequencies [Hz]' ) 
    odmr_line_width = Property( trait=Float, depends_on='fit_parameters', label='line width [Hz]'  ) 
    odmr_contrast = Property( trait=Float, depends_on='fit_parameters', label='contrast [%]'  ) 

    frequency = Array()
    counts = Array(comparison_mode=NO_COMPARE)

    counts_matrix = Array()

    line_plot_text = Instance( PlotLabel, factory=PlotLabel, kw={'text':'', 'hjustify':'left', 'vjustify':'bottom', 'position':[64,32]})
    line_plot_data = Instance( ArrayPlotData )
    line_plot = Instance( Plot )
    matrix_plot_data = Instance( ArrayPlotData )
    matrix_plot = Instance( Plot )

    def __init__(self):
        super(ODMR, self).__init__()
        self.on_trait_change(self._update_line_plot_index,      'frequency',            dispatch='ui')
        self.on_trait_change(self._update_line_plot_value,      'counts',               dispatch='ui')
        self.on_trait_change(self._update_matrix_plot_value,    'counts_matrix',        dispatch='ui')
        self.on_trait_change(self._update_matrix_plot_index,    'n_lines,frequency',    dispatch='ui')
        self.on_trait_change(self._update_line_plot_fit,        'fit_parameters',       dispatch='ui')

    def submit(self):
        """Submit the job to the JobManager."""
        self.keep_data = False
        ManagedJob.submit(self)

    def resubmit(self):
        """Submit the job to the JobManager."""
        self.keep_data = True
        ManagedJob.submit(self)

    def _resubmit_button_fired(self):
        """React to start button. Submit the Job."""
        self.resubmit() 

    def _line_plot_data_default(self):
        return ArrayPlotData(frequency=self.frequency, counts=self.counts)
    def _line_plot_default(self):
        plot = Plot(self.line_plot_data, padding=8, padding_left=64, padding_bottom=32)
        plot.plot(('frequency','counts'), style='line', color='blue')
        plot.index_axis.title = 'Frequency [Hz]'
        plot.value_axis.title = 'Fluorescence counts'
        plot.overlays.append(self.line_plot_text)
        return plot

    def _counts_matrix_default(self):
        return np.zeros( (self.n_lines, len(self.frequency)) )

    def _matrix_plot_data_default(self):
        return ArrayPlotData(image=self.counts_matrix)
    def _matrix_plot_default(self):
        plot = Plot(self.matrix_plot_data, padding=8, padding_left=64, padding_bottom=32)
        plot.index_axis.title = 'Frequency [MHz]'
        plot.value_axis.title = 'line #'
        plot.img_plot('image',
                      xbounds=(self.frequency[0],self.frequency[-1]),
                      ybounds=(0,self.n_lines),
                      colormap=Spectral)
        return plot

    def _frequency_default(self):
        return np.arange(self.frequency_begin, self.frequency_end+self.frequency_delta, self.frequency_delta)

    def _counts_default(self):
        return np.zeros(self.frequency.shape)

    def _update_line_plot_index(self):
        self.line_plot_data.set_data('frequency', self.frequency)
        self.counts_matrix = self._counts_matrix_default()

    def _update_line_plot_value(self):
        self.line_plot_data.set_data('counts', self.counts)

    def _perform_fit_changed(self):
        plot = self.line_plot
        if self.perform_fit is False:
            plot.delplot('fit')
            self.line_plot_text.visible=False
        else:
            plot.plot(('frequency','fit'), style='line', color='red', name='fit')
            self.line_plot_text.visible=True
        plot.request_redraw()

    def _update_matrix_plot_value(self):
        self.matrix_plot_data.set_data('image', self.counts_matrix)

    def _update_matrix_plot_index(self):
        if self.n_lines > self.counts_matrix.shape[0]:
            self.counts_matrix = np.vstack( (self.counts_matrix,np.zeros((self.n_lines-self.counts_matrix.shape[0], self.counts_matrix.shape[1]))) )
        else:
            self.counts_matrix = self.counts_matrix[:self.n_lines]        
        self.matrix_plot.components[0].index.set_data((self.frequency[0], self.frequency[-1]),(0.0,float(self.n_lines)))

    def apply_parameters(self):
        """Apply the current parameters and decide whether to keep previous data."""
        frequency = np.arange(self.frequency_begin, self.frequency_end+self.frequency_delta, self.frequency_delta)

        if not self.keep_data or np.any(frequency != self.frequency):
            self.counts = np.zeros(frequency.shape)            

        self.frequency = frequency
        self.keep_data = True # when job manager stops and starts the job, data should be kept. Only new submission should clear data.

    def _run(self):
        
        try:
            self.state='run'
            self.apply_parameters()

            # if pulsed, turn on sequence
            if self.pulsed:
                ha.PulseGenerator().Sequence( 100 * [ (['laser','aom'],self.laser), ([],self.wait), (['mw'],self.t_pi) ] )
            else:
                ha.PulseGenerator().Open()

            n = len(self.frequency)

            """
            ha.Microwave().setOutput( self.power, np.append(self.frequency,self.frequency[0]), self.seconds_per_point)
            self._prepareCounter(n)
            """
            ha.Microwave().setPower(self.power)
            ha.Microwave().initSweep( self.frequency, self.power*np.ones(self.frequency.shape))
            ha.Counter().configure(n, self.seconds_per_point, DutyCycle=0.8)
            time.sleep(0.5)

            self.i = 0
            self.elapsed_time = 0
            self.progress = 0
            start_time = time.time()
            while self.i < self.run_time / ( self.seconds_per_point * n ):
                if threading.currentThread().stop_request.isSet():
                    break
                ha.Microwave().resetListPos()
                counts = ha.Counter().run()
                
                """
                ha.Microwave().doSweep()
                
                timeout = 3.
                start_time = time.time()
                while not self._count_between_markers.ready():
                    time.sleep(0.1)
                    if time.time() - start_time > timeout:
                        print "count between markers timeout in ODMR"
                        break
                        
                counts = self._count_between_markers.getData(0)
                """
                self.counts += counts
                self.counts_matrix = np.vstack( (counts, self.counts_matrix[:-1,:]) )
                """
                self._count_between_markers.clean()
                """
                self.elapsed_time = int(time.time() - start_time)
                self.progress = int( 100 * float(self.i) / self.run_time * self.seconds_per_point * n  )            
                self.i=self.i+1
    
            ha.Microwave().setOutput( None, self.frequency_begin)
            ha.PulseGenerator().Light()
            ha.Counter().clear()
            
        finally:
            self.state = 'idle'

    @cached_property
    def _get_fit_parameters(self):
        if self.perform_fit:
            N = self.number_of_dips 
            if N != 'auto':
                N = int(N)
            try:
                p = fitting.fit_multiple_lorentzians(self.frequency,self.counts,N)
            except Exception:
                logging.getLogger().debug('fit did not work out', exc_info=True)
                #p = np.array( (np.NaN, np.NaN, np.NaN, np.NaN) )
                p = np.array( (0., 0., 0., 0.) )
        else:
            #p = np.array( (np.NaN, np.NaN, np.NaN, np.NaN) )
            p = np.array( (0., 0., 0., 0.) )
        return p

    def _update_line_plot_fit(self):
        #if self.fit_parameters[0] is not np.NaN:
        if not self.fit_parameters[0] == 0.:
            self.line_plot_data.set_data('fit', fitting.NLorentzians(*self.fit_parameters)(self.frequency))
            p = self.fit_parameters
            f = p[1::3]
            w = p[2::3]
            c = -100*p[3::3] / (np.pi*w*p[0])
            s = ''
            for i, fi in enumerate(f):
                s += 'dip %i: %.6e Hz +- %.3e Hz, %.1f%%\n'%(i+1, fi, w[i], c[i])
            self.line_plot_text.text = s
            

    @cached_property
    def _get_odmr_frequencies(self):
        return self.fit_parameters[1::3]

    @cached_property
    def _get_odmr_line_width(self):
        return self.fit_parameters[2::3].mean()

    @cached_property
    def _get_odmr_contrast(self):
        return -100*self.fit_parameters[3::3].mean()/(np.pi*self.fit_parameters[2::3].mean()*self.fit_parameters[0])

    def save_line_plot(self, filename):
        self.save_figure(self.line_plot, filename)

    def save_matrix_plot(self, filename):
        self.save_figure(self.matrix_plot, filename)
    
    def save_all(self, filename):
        self.save_line_plot(filename+'_ODMR_Line_Plot.png')
        self.save_matrix_plot(filename+'_ODMR_Matrix_Plot.png')
        self.save(filename+'_ODMR.pys')

    traits_view = View(VGroup(HGroup(Item('submit_button',   show_label=False),
                                     Item('remove_button',   show_label=False),
                                     Item('resubmit_button', show_label=False),
                                     Item('priority'), Item('state', style='readonly'),
                                     ),
                              Group(VGroup(HGroup(Item('power', width=-40),
                                                  Item('frequency_begin', width=-80, editor=TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x:'%e'%x)),
                                                  Item('frequency_end', width=-80, editor=TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x:'%e'%x)),
                                                  Item('frequency_delta', width=-80, editor=TextEditor(auto_set=False, enter_set=True, evaluate=float, format_func=lambda x:'%e'%x)),
                                                  Item('pulsed'),
                                                  Item('t_pi', width=-80),
                                                  ),
                                           HGroup(Item('run_time'),
                                                  Item('perform_fit'),
                                                  Item('number_of_dips'),
                                                  Item('n_lines'),
                                                  ),
                                           HGroup(Item('elapsed_time', style='readonly'),
                                                  Item('progress', style='readonly'),
                                                  Item('odmr_contrast', style='readonly', editor=TextEditor(evaluate=float, format_func=lambda x:'%.1f'%x)),
                                                  Item('odmr_line_width', style='readonly', editor=TextEditor(evaluate=float, format_func=lambda x:'%.3e'%x)),
                                                  Item('odmr_frequencies', style='readonly'),
                                                  ),
                                           label='parameters'
                                           ),
                                    VGroup(HGroup(Item('seconds_per_point'),),
                                           HGroup(Item('laser', width=40),
                                                  Item('wait', width=40),
                                                  ),
                                           label='settings'
                                           ),
                                    layout='tabbed',
                                    ),
                              VSplit(Item('matrix_plot', editor=ComponentEditor(), show_label=False, resizable=True),
                                     Item('line_plot', editor=ComponentEditor(), show_label=False, resizable=True),
                                     ),
                              ),
                       menubar = MenuBar(Menu(Action(action='saveLinePlot', name='SaveLinePlot (.png)'),
                                              Action(action='saveMatrixPlot', name='SaveMatrixPlot (.png)'),
                                              Action(action='save', name='Save (.pyd or .pys)'),
                                              Action(action='saveAll', name='Save All (.png+.pys)'),
                                              Action(action='export', name='Export as Ascii (.asc)'),
                                              Action(action='load', name='Load'),
                                              Action(action='_on_close', name='Quit'),
                                              name='File')),
                       title='ODMR', width=900, height=800, buttons=[], resizable=True, handler=ODMRHandler
                       )

    get_set_items=['power', 'frequency_delta', 'frequency_begin', 'frequency_end', 'frequency', 'seconds_per_point', 'run_time', 'n_lines',
                   'pulsed', 't_pi', 'laser', 'wait',
                   'counts', 'counts_matrix', 'elapsed_time', 'progress',
                   'number_of_dips', 'perform_fit',
                   'fit_parameters', 'odmr_contrast', 'odmr_line_width', 'odmr_frequencies', '__doc__' ]
    get_set_order=['frequency','n_lines']


if __name__ == '__main__':

    logging.getLogger().addHandler(logging.StreamHandler())
    logging.getLogger().setLevel(logging.DEBUG)
    logging.getLogger().info('Starting logger.')

    o = ODMR()
    o.edit_traits()
    